原始题目：剑指 Offer 41. 数据流中的中位数 (opens new window)

解题思路：

中位数又叫中值，是按顺序排列的一组数据中居于中间位置的数。

因为需要顺序排列，所以我们需要对数据流中的数字进行排序。

自定义规则

以中位数为界限，我们将小于中位数的数字成为较小数，大于中位数的数字成为较大数。较小数的那部分称为较小部分，较大数的部分称为较大部分。然后规定添加过程中，要保持较大部分的个数大于等于较小部分的个数。

我们在计算中位数的时候，需要分为两种情况，元素总个数为奇数和偶数：

• 元素总数为奇数时，需要拿到较大部分中的最小数作为中位数返回；
• 元素总数为偶数时，需要拿到较大部分的最小数和叫小部分的最大数的平均数返回；

数据结构选型

接下来就是思考如何存储较小部分，然后可以快速拿到较小部分的最大数，以及如何存储较大部分，然后可以快速拿到较大部分的最小数？

• 可以使用大根堆来存储较小部分，因为大根堆的根是所有堆内元素最大的数字；
• 可以使用小根堆来存储较大部分，因为小根堆的根是所有堆内元素最小的数字；

思路梳理

1. 将数据流的数据，分为较大部分（用小根堆 $A$ 存储）和较小部分（用大根堆 $B$ 存储）。
2. 添加数字 $x$ 的时候
   1. 如果此时总数为偶数，那么根据前面的规定，添加后 $A$ 要比 $B$ 多一个数字，先将 $x$ 添加到 $B$ 中，然后将 $B$ 的根添加到 $A$ 中。
   2. 如果此时总数为奇数，此时 $A.size > B.size$，那么为了维持 $A$ 和 $B$ 的大小平衡，添加后 $A$ 和 $B$ 的个数要想等，先将 $x$ 添加到 $A$ 中，然后将 $A$ 的根添加到 $B$ 中。
   3. 总数加一。
3. 求中位数时
   1. 如果总数为奇数，那么直接返回 $A$ 的根。
   2. 如果总数为偶数，那么返回 $A$ 和 $B$ 的根的平均数。

代码：

class MedianFinder {
    /**
     * 小顶堆存放的是数据流排序后的后半部分，即较大的一半
     */
    private final PriorityQueue<Integer> minHeap;
    /**
     * 大顶堆存放的是数据流排序后的前半部分，即较小的一半
     */
    private final PriorityQueue<Integer> maxHeap;

    private int total = 0;

    /**
     * initialize your data structure here.
     */
    public MedianFinder() {
        minHeap = new PriorityQueue<>(Comparator.naturalOrder());
        maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
    }

    /**
     * 添加的过程中，要保持小顶堆的元素个数大于等于大顶堆的个数，
     * 即总元素个数为奇数的时候，小顶堆要比大顶堆多一个元素
     */
    public void addNum(int num) {
        // 如果是偶数，则入小顶堆，但是不能直接入
        // 原因是，这个数字可能属于较小的一半，所以需要先入大顶堆，再把大顶堆的根弹出
        if ((total & 1) == 0) {
            maxHeap.offer(num);
            minHeap.offer(maxHeap.poll());
        } else {
            // 如果是奇数，则入大顶堆，但是不能直接入
            // 原因是，这个数字可能属于较大的一半，所以需要先入小顶堆，再把小顶堆的根弹出
            minHeap.offer(num);
            maxHeap.offer(minHeap.poll());
        }
        total++;
    }

    public double findMedian() {
        // 偶数则返回中间两个元素的和的一半
        if ((total & 1) == 0) {
            return (double) (minHeap.peek() + maxHeap.peek()) / 2;
        } else {
            // 奇数则返回小顶堆的根
            return (double) minHeap.peek();
        }
    }
}

复杂度分析

• 时间复杂度

  • 添加数字 $O(logN)$：堆的弹出和压入使用 $O(logN)$ 的时间；
  • 求中位数$O(1)$：获取堆顶元素使用 $O(1)$ 时间。

• 空间复杂度$O(N)$：其中 $N$ 为数据流中的元素数量，小顶堆 $A$ 和大顶堆 $B$ 最多同时保存 $N$ 个元素。